import sys
sys.path.append('/usr/lib/python3.8/site-packages')
from flask import Flask, request, jsonify, redirect, send_from_directory,send_file
import requests
from transformers import WhisperProcessor, WhisperForConditionalGeneration,AutoModelForSpeechSeq2Seq,QuantoConfig,AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import numpy as np
import torch
import io
import os
import base64
import soundfile as sf
import edge_tts
import asyncio


app = Flask(__name__)
has_initialized = False
# 确保上传文件夹存在
if not os.path.exists('/tmp/uploads'):
    os.makedirs('/tmp/uploads')

# 使用 app.before_first_request 装饰器来初始化模型

def load_model():
    device ='cpu'
    model_id = "xmzhu/whisper-tiny-zh"
    quanto_config = QuantoConfig(weights="int8")

    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        model_id,
        # torch_dtype=torch.float32,
        device_map="cpu",
        # low_cpu_mem_usage=True,
        quantization_config=quanto_config
    )

    processor = AutoProcessor.from_pretrained(model_id)
    #freeze(model)
    model = torch.compile(model)
    
    app.transcriber = pipeline(
        "automatic-speech-recognition",
        model=model,
        tokenizer=processor.tokenizer,
        feature_extractor=processor.feature_extractor,
        # max_new_tokens=128,
        # torch_dtype=torch.float32,
    )
    app.transcriber.model.config.forced_decoder_ids = app.transcriber.tokenizer.get_decoder_prompt_ids(language="zh", task="transcribe")
    

def initialize():
    load_model()
		#初始化函数   我们需要初始执行的函数
    print('hello world.................654646546465..')
    pass
    

@app.before_request
def first_request():
    global has_initialized
    if not has_initialized:
        initialize()
        has_initialized = True
        print('hello world...................')


@app.route('/v1/audio/transcriptions', methods=['POST'])
def upload_audio():
    if request.method == 'POST':
        # 获取上传的录音文件
        audio_file = request.files['audio_file']
        if audio_file is None:
            return jsonify({'error': 'Audio file not provided'}), 400

        # 保存音频文件到本地
        file_path = os.path.join('/tmp/uploads', audio_file.filename)
        audio_file.save(file_path)

        # 对音频文件进行转录
        transcription = app.transcriber(file_path)
        print(transcription)

        # 删除上传的音频文件
        os.remove(file_path)

        # 返回转录的文字
        return jsonify({'transcription': transcription}), 200

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8060)